import unittest

import pandas as pd

from download_additional import get_code, validate_data


class TestStockDataDownloader(unittest.TestCase):
    def test_get_code(self):
        codes = get_code()
        self.assertIn('上证综指', codes)
        self.assertEqual(codes['上证综指'], 'sh000001')

    def test_validate_data(self):
        df = pd.DataFrame({'date': ['2021-01-01'], 'value': [100]})
        self.assertIsNone(validate_data(df))
        with self.assertRaises(ValueError):
            validate_data(pd.DataFrame())


if __name__ == '__main__':
    unittest.main()
